package com.atguigu.day08;

import org.apache.flink.api.common.typeutils.base.StringSerializer;
import org.apache.flink.api.common.typeutils.base.VoidSerializer;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.sink.TwoPhaseCommitSinkFunction;
import org.apache.flink.streaming.api.functions.source.SourceFunction;

import java.io.BufferedWriter;
import java.nio.file.Files;
import java.nio.file.Paths;

public class Example3 {
    public static void main(String[] args) {
        var env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);

        env.enableCheckpointing(10 * 1000L);

        env
                .addSource(new SourceFunction<Long>() {
                    private boolean running = true;
                    private long count = 1L;
                    @Override
                    public void run(SourceContext<Long> ctx) throws Exception {
                        while (running) {
                            ctx.collect(count);
                            count++;

                            Thread.sleep(1000L);
                        }
                    }

                    @Override
                    public void cancel() {
                        running = false;
                    }
                })
                .addSink(new TransactionalFileSink());

        try {
            env.execute();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    // 三个泛型：输入数据的泛型，事务的泛型，上下文的泛型
    public static class TransactionalFileSink extends TwoPhaseCommitSinkFunction<Long, String, Void> {
        // 声明一个缓冲区
        private BufferedWriter bufferedWriter;

        public TransactionalFileSink() {
            super(StringSerializer.INSTANCE, VoidSerializer.INSTANCE);
        }

        // 开启事务时触发执行
        @Override
        protected String beginTransaction() throws Exception {
            long timeNow = System.currentTimeMillis();
            int taskIdx = getRuntimeContext().getIndexOfThisSubtask();
            // 创建临时文件名
            String transaction = timeNow + "-" + taskIdx;
            // 获取临时文件的绝对路径
            var tempFilePath = Paths.get("/home/yuantuzhi/flinktutorial0926/src/main/resources/tempfile/" + transaction);
            // 创建文件
            Files.createFile(tempFilePath);
            // 将缓冲区指定为临时文件的缓冲区
            bufferedWriter = Files.newBufferedWriter(tempFilePath);
            // 返回事务
            return transaction;
        }

        // 每来一条数据触发一次调用
        @Override
        protected void invoke(String transaction, Long value, Context context) throws Exception {
            // 将到达的数据写入临时文件的缓冲区中
            bufferedWriter.write(value + "\n");
        }

        // 两阶段提交的第一阶段：预提交
        @Override
        protected void preCommit(String transaction) throws Exception {
            // 将缓冲区中的数据刷入临时文件
            bufferedWriter.flush();
            // 关闭缓冲区
            bufferedWriter.close();
        }

        // 当sink算子接收到作业管理器发来的检查点完成的通知时，触发调用
        // 两阶段提交的第二阶段：正式提交
        @Override
        protected void commit(String transaction) {
            // 获取临时文件的绝对路径
            var tempFilePath = Paths.get("/home/yuantuzhi/flinktutorial0926/src/main/resources/tempfile/" + transaction);

            if (Files.exists(tempFilePath)) {
                try {
                    // 正式文件的绝对路径
                    var commitFilePath = Paths.get("/home/yuantuzhi/flinktutorial0926/src/main/resources/targetfile/" + transaction);
                    // 将临时文件的绝对路径更名为正式文件的绝对路径
                    Files.move(tempFilePath, commitFilePath);
                } catch (Exception e) {
                    e.printStackTrace();
                }
            } else {
                System.out.println("临时文件不存在");
            }
        }

        // 回滚时触发调用
        @Override
        protected void abort(String transaction) {
            var tempFilePath = Paths.get("/home/yuantuzhi/flinktutorial0926/src/main/resources/tempfile/" + transaction);
            if (Files.exists(tempFilePath)) {
                try {
                    Files.delete(tempFilePath);
                } catch (Exception e) {
                    e.printStackTrace();
                }
            } else {
                System.out.println("临时文件不存在");
            }
        }
    }
}
